from langgraph_supervisor import create_supervisor
from langgraph.prebuilt import create_react_agent

from common.ali_qwen import get_model

model = get_model()
#为专业智能体定义工具
def add(a:float,b:float) ->float:
    """两个数相加"""
    return a+b

def multiply(a:float,b:float) ->float:
    """两个数相乘"""
    return a*b

def web_search(query:str) -> str:
    """在网络上搜索信息"""
    return (
        "以下是 FAANG 公司2024年的员工人数：\n"
        "1:**Facebook** 67317 个员工\n",
        "2:**Apple** 16400 个员工\n",
        "3:**Amazon** 55100 个员工\n",
        "4:**Netflix** 14000个员工\n",
        "5:**Google** 181236个员工\n",
    )

math_agent = create_react_agent(
    model=model,
    tools=[add,multiply],
    name="math_exper",
    prompt="你是一个数学家，始终使用一个工具"
)

research_agent = create_react_agent(
    model=model,
    tools=[web_search],
    name="research_exper",
    prompt="你是一个世界一流的研究专家，可以进行网络搜索，不要做任何的运算。"
)

#使用create supervisor创建一个主管智能体
workflow = create_supervisor(
    (research_agent,math_agent),
    model=model,
    prompt="你是一个团队主管，管理一个研究专家和一位数学专家.研究专家能利用网络搜索工具进行查询。"
)
app = workflow.compile()
result = app.invoke({
    "messages":[
        {
            "role":"user",
            "content":"2024年FAANG公司的总员工人数多少？"
        }
    ]
})
print(result["messages"][-1].content)